package websocket

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path"
	"strconv"
	"strings"

	"time"
	"tuiwen-go/internal/api"
	"tuiwen-go/internal/config"
	"tuiwen-go/internal/db"
	"tuiwen-go/internal/utils"
)

// taskImgHandler handles messages of type 4.
func taskImgByComfyuiHandler(c *Client, data json.RawMessage) {
	log.Printf("Client %s: Handling Image Task, Data: %s", c.ID, string(data))

	var taskData struct {
		ProjectID string `json:"project_id"`
		Id        string `json:"id"`
	} // Assuming the data contains project_id and project_name

	if err := json.Unmarshal(data, &taskData); err != nil {
		log.Printf("Client %s: Error unmarshalling image task data: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "Invalid image task data format"})
		return
	}

	if taskData.ProjectID == "" {
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "项目ID不能为空"})
		return
	}

	// Fetch text entries that need images generated
	sqlQuery := "SELECT t.id, t.prompt, image,r.imgPath imgPath FROM text t LEFT JOIN role r on t.role_ids = r.id WHERE project_id = ? AND content IS NOT NULL"
	args := []interface{}{taskData.ProjectID}
	if taskData.Id != "" {
		sqlQuery += " AND t.id = ?"
		args = append(args, taskData.Id)
	} else {
		sqlQuery += " AND (image IS NULL OR image = '')"
	}
	rows, err := db.Query(sqlQuery, args...)
	if err != nil {
		log.Printf("Client %s: Error querying text for image task: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "查询待生成图片文本失败"})
		return
	}
	defer rows.Close()

	var textsToImage []struct {
		ID      int
		Prompt  string
		Image   string // Add Image field
		imgPath string
	}

	for rows.Next() {
		var id int
		var prompt sql.NullString
		var image sql.NullString // Scan the image column
		var imgPath sql.NullString
		err := rows.Scan(&id, &prompt, &image, &imgPath)
		if err != nil {
			log.Printf("Client %s: Error scanning text row for image task: %v", c.ID, err)
			continue
		}

		if prompt.Valid {
			textsToImage = append(textsToImage, struct {
				ID      int
				Prompt  string
				Image   string
				imgPath string
			}{ID: id, Prompt: prompt.String, Image: image.String, imgPath: imgPath.String})
		}
	}

	if err = rows.Err(); err != nil {
		log.Printf("Client %s: Error iterating text rows for image task: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "处理待生成图片文本失败"})
		return
	}

	if len(textsToImage) == 0 {
		sendMessage(c, api.WxJSONResponse{Code: 0, Msg: "图片已全部生成"})
		return // No texts to generate images for
	}
	var lastId int
	if taskData.Id != "" {
		// 将字符串ID转换为整数并减1
		id, err := strconv.Atoi(taskData.Id)
		if err != nil {
			log.Printf("Error converting ID to integer: %v", err)
			lastId = 0
		} else {
			lastId = id - 1
		}
	}

	// Process each text entry to generate image
	for _, item := range textsToImage {
		// Placeholder for image generation logic
		// This requires calling an external service or a local model.
		// For now, let's use dummy values.
		// In a real implementation, integrate an image generation function.
		generatedImage, image_comfyui, err := generateAndSaveImageByComfyui(item.Prompt, item.imgPath, taskData.ProjectID, lastId)
		if err != nil {
			log.Printf("Client %s: Error generating and saving image for text ID %d: %v", c.ID, item.ID, err)
			// Continue processing other items, but log the error
			sendMessage(c, api.WxJSONResponse{Code: 1, Msg: fmt.Sprintf("图片生成失败 (ID: %d): %v", item.ID, err)})
		} else {

			// Update the database with the generated image path
			sqlUpdate := "UPDATE text SET image = ?,image_comfyui = ? WHERE project_id = ? AND id = ?"
			_, err := db.Exec(sqlUpdate, generatedImage, image_comfyui, taskData.ProjectID, item.ID)
			if err != nil {
				log.Printf("Client %s: Error updating image for text ID %d: %v", c.ID, item.ID, err)
				// Continue processing other items, but log the error
			} else {

				// 旧图片重命名，格式为：新图片名_旧图片名.png
				if item.Image != "" {
					oldImagePath := strings.ReplaceAll(item.Image, "/public", "./public")
					newImageFilename := path.Base(generatedImage)
					oldImageFilename := path.Base(oldImagePath)
					newOldImageName := strings.TrimSuffix(newImageFilename, path.Ext(newImageFilename)) + "_" + oldImageFilename
					oldImageDir := path.Dir(oldImagePath)
					newOldImagePath := path.Join(oldImageDir, newOldImageName)

					log.Printf("Client %s: Renaming old image file from %s to %s", c.ID, oldImagePath, newOldImagePath)
					err := os.Rename(oldImagePath, newOldImagePath)
					if err != nil {
						log.Printf("Client %s: Failed to rename old image file %s to %s: %v", c.ID, oldImagePath, newOldImagePath, err)
					}
				}
				// Send a message to the client with the updated image URL
				sendMessage(c, api.WxJSONResponse{
					Code: 1, // Or a specific code for image update
					Msg:  fmt.Sprintf("配图生成成功 (ID: %d)", item.ID),
					Data: map[string]interface{}{
						"id":         item.ID,
						"image_path": generatedImage,
					},
				})
			}
		}
	}

	// A final success message after attempting all updates (optional, as individual messages are sent)
	// sendMessage(c, api.WxJSONResponse{Code: 0, Msg: "图片生成任务完成"})
}

func generateAndSaveImageByComfyui(prompt string, imagePath string, projectID string, lastId int) (string, string, error) {
	log.Printf("Generating image for prompt: %s (Project ID: %s)", prompt, projectID)

	// 1. Load the ComfyUI API prompt template
	promptTemplatePath, ok := config.GetSetting("comfyUI_template_path")
	if !ok || promptTemplatePath == "" {
		promptTemplatePath = "templates/comfyuiApiJson/kentext加速版.json"
	}

	exeutableDir, err := os.Executable()
	if err != nil {
		return "", "", fmt.Errorf("error getting executable directory: %v", err)
	}
	jsonFilePath := path.Join(path.Dir(exeutableDir), promptTemplatePath)

	jsonFile, err := os.ReadFile(jsonFilePath)
	if err != nil {
		return "", "", fmt.Errorf("error reading comfyui api json file: %v", err)
	}

	var promptData map[string]interface{}
	if err := json.Unmarshal(jsonFile, &promptData); err != nil {
		return "", "", fmt.Errorf("error unmarshalling comfyui api json: %v", err)
	}

	// 2. Modify the prompt with dynamic values
	// Example: Find the node for text prompt and update it.
	// This needs to be adapted based on the actual structure of your JSON template.
	promptNodeID, ok := config.GetSetting("comfyUI_prompt_node_id")
	if !ok {
		promptNodeID = "190"
	}

	if node, ok := promptData[promptNodeID].(map[string]interface{}); ok {
		if inputs, ok := node["inputs"].(map[string]interface{}); ok {
			inputs["text"] = prompt
		} else {
			return "", "", fmt.Errorf("invalid comfyui template structure: 'inputs' not found or not a map")
		}
	} else {
		return "", "", fmt.Errorf("invalid comfyui template structure: node '%s' not found or not a map", promptNodeID)
	}

	noise_seed := int64(time.Now().UnixNano())

	if node, ok := promptData["186"].(map[string]interface{}); ok {
		if inputs, ok := node["inputs"].(map[string]interface{}); ok {
			inputs["seed"] = noise_seed
		} else {
			return "", "", fmt.Errorf("invalid comfyui template structure: 'inputs' not found or not a map")
		}
	} else {
		return "", "", fmt.Errorf("invalid comfyui template structure: node '%s' not found or not a map", promptNodeID)
	}

	comfyUI_imgUpload_node_id, ok := config.GetSetting("comfyUI_imgUpload_node_id")
	if !ok {
		comfyUI_imgUpload_node_id = "191"
	}
	if imagePath == "" {
		// 根据lastId 查询text image_comfyui
		err := db.QueryRow("SELECT image_comfyui FROM text WHERE id = ?", lastId).Scan(&imagePath)
		if err != nil {
			return "", "", fmt.Errorf("error querying image_comfyui: %v", err)
		}
	}
	imagePath = strings.ReplaceAll(imagePath, "\\", "/")
	if strings.HasPrefix(imagePath, "/public") {
		// 获取当前工作目录
		currentDir, err := os.Getwd()
		if err != nil {
			fmt.Println("获取当前工作目录出错:", err)
		}
		imagePath = currentDir + imagePath
	}

	if node, ok := promptData[comfyUI_imgUpload_node_id].(map[string]interface{}); ok {
		if inputs, ok := node["inputs"].(map[string]interface{}); ok {
			inputs["image"] = imagePath
		} else {
			return "", "", fmt.Errorf("invalid comfyui template structure: 'inputs' not found or not a map")
		}
	} else {
		return "", "", fmt.Errorf("invalid comfyui template structure: node '%s' not found or not a map", comfyUI_imgUpload_node_id)
	}

	comfyUI_image_node_id, ok := config.GetSetting("comfyUI_image_node_id")
	if !ok {
		comfyUI_image_node_id = "27"
	}
	width, ok := config.GetSetting("image_width")
	if !ok {
		width = "512"
	}
	height, ok := config.GetSetting("image_height")
	if !ok {
		height = "512"
	}

	if node, ok := promptData[comfyUI_image_node_id].(map[string]interface{}); ok {
		if inputs, ok := node["inputs"].(map[string]interface{}); ok {
			inputs["width"] = width
			inputs["height"] = height
		} else {
			return "", "", fmt.Errorf("invalid comfyui template structure: 'inputs' not found or not a map")
		}
	} else {
		return "", "", fmt.Errorf("invalid comfyui template structure: node '%s' not found or not a map", comfyUI_image_node_id)
	}
	// 3. Call the ComfyUI API
	outputImages, image_comfyui, err := utils.GetImages(promptData)
	if err != nil {
		return "", "", fmt.Errorf("error getting images from comfyui: %v", err)
	}

	// 4. Find the output image from the 'SaveImage' node
	var imageData []byte
	found := false
	for _, nodeOutput := range outputImages {
		if len(nodeOutput) > 0 {
			imageData = nodeOutput[0] // Assuming the first image is the one we want
			found = true
			break
		}
	}

	if !found || len(imageData) == 0 {
		return "", "", fmt.Errorf("no images returned from ComfyUI")
	}

	// 5. Save the image
	dir := fmt.Sprintf("public/images/%s", projectID)
	saveDir := path.Join(path.Dir(exeutableDir), dir)

	if _, err := os.Stat(saveDir); os.IsNotExist(err) {
		err = os.MkdirAll(saveDir, 0755)
		if err != nil {
			return "", "", fmt.Errorf("error creating directory %s: %v", saveDir, err)
		}
	}

	filename := fmt.Sprintf("%d.png", time.Now().UnixNano())
	savePath := path.Join(saveDir, filename)

	err = os.WriteFile(savePath, imageData, 0644)
	if err != nil {
		return "", "", fmt.Errorf("error saving image file %s: %v", savePath, err)
	}

	publicPath := path.Join("images", projectID, filename)
	publicPath = strings.ReplaceAll(publicPath, "\\", "/")

	log.Printf("Image saved to %s, public path: %s", savePath, publicPath)

	return "/public/" + publicPath, image_comfyui + " [output]", nil
}

// 微调指定图片
func taskImgByComfyuiContextHandler(c *Client, data json.RawMessage) {
	log.Printf("Client %s: Handling Image Task, Data: %s", c.ID, string(data))
	// 1. Unmarshal the data
	var taskData struct {
		ProjectID string `json:"project_id"`
		Id        string `json:"id"`
		Prompt    string `json:"prompt"`
		ImgPath   string `json:"imgPath"`
	} // Assuming the data contains project_id and project_name

	if err := json.Unmarshal(data, &taskData); err != nil {
		log.Printf("Client %s: Error unmarshalling image task data: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "Invalid image task data format"})
		return
	}

	if taskData.ProjectID == "" {
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "项目ID不能为空"})
		return
	}
	id, err := strconv.Atoi(taskData.Id)
	if err != nil {
		log.Printf("Client %s: Error parsing text ID %s: %v", c.ID, taskData.Id, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "Invalid text ID format"})
		return
	}

	generatedImage, image_comfyui, err := generateAndSaveImageByComfyui(taskData.Prompt, taskData.ImgPath, taskData.ProjectID, id)
	if err != nil {
		log.Printf("Client %s: Error generating and saving image for text ID %s: %v", c.ID, taskData.Id, err)
		// Continue processing other items, but log the error
		sendMessage(c, api.WxJSONResponse{Code: 1, Msg: fmt.Sprintf("图片生成失败 (ID: %s): %v", taskData.Id, err)})
	} else {

		// Update the database with the generated image path
		sqlUpdate := "UPDATE text SET image = ?, image_comfyui = ? WHERE project_id = ? AND id = ?"
		_, err := db.Exec(sqlUpdate, generatedImage, image_comfyui, taskData.ProjectID, taskData.Id)
		if err != nil {
			log.Printf("Client %s: Error updating image for text ID %s: %v", c.ID, taskData.Id, err)
			// Continue processing other items, but log the error
		} else {

			// 旧图片重命名，格式为：新图片名_旧图片名.png

			oldImagePath := strings.ReplaceAll(taskData.ImgPath, "/public", "./public")
			newImageFilename := path.Base(generatedImage)
			oldImageFilename := path.Base(oldImagePath)
			newOldImageName := strings.TrimSuffix(newImageFilename, path.Ext(newImageFilename)) + "_" + oldImageFilename
			oldImageDir := path.Dir(oldImagePath)
			newOldImagePath := path.Join(oldImageDir, newOldImageName)

			log.Printf("Client %s: Renaming old image file from %s to %s", c.ID, oldImagePath, newOldImagePath)
			err := os.Rename(oldImagePath, newOldImagePath)
			if err != nil {
				log.Printf("Client %s: Failed to rename old image file %s to %s: %v", c.ID, oldImagePath, newOldImagePath, err)
			}

			// Send a message to the client with the updated image URL
			sendMessage(c, api.WxJSONResponse{
				Code: 1, // Or a specific code for image update
				Msg:  fmt.Sprintf("配图生成成功 (ID: %s)", taskData.Id),
				Data: map[string]interface{}{
					"id":            taskData.Id,
					"image_path":    generatedImage,
					"image_comfyui": image_comfyui,
				},
			})
		}
	}
}
